#!/usr/bin/env python
# encoding: utf-8
'''
@author: Justin Ruan
@license: 
@contact: ruanjun@whut.edu.cn
@time: 2020-01-05
@desc:
'''
import networkx as nx
import csv
import util
import numpy as np
from bidict import bidict

PROJECT_ROOT = util.get_project_root()


class Combiner(object):
    def __init__(self):
        pass

    def loading_saved_graphs(self, filename):
        '''

        :return:
        '''
        filename = "{}//results//graph//{}".format(PROJECT_ROOT, filename)
        data = np.load(filename, allow_pickle=True)
        graphs = data["g"].item()
        return graphs

    def assembling_network(self, graphs):
        net = nx.Graph()
        dict_code_entity = bidict()
        edges = []
        for g in graphs.values():
            for id, node in g.nodes(data=True):
                label = node['type']
                entity = node['name']
                if label not in ['A', "C"]:
                    if net.has_node(id):
                        weight = net.nodes[id]['weight']
                        net.nodes[id]['weight'] = weight + 1
                    else:
                        net.add_node(id, type=label, name=entity, weight=1)
                        dict_code_entity[id] = entity

            for start, end in g.edges():
                if net.has_node(start) and net.has_node(end):
                    if net.has_edge(start, end):
                        weight = net.edges[start, end]['weight']
                        net.edges[start, end]['weight'] = weight + 1
                    else:
                        net.add_edge(start, end, weight=1)

        return net, dict_code_entity


    @staticmethod
    def export_csv(filename):
        '''
        将存盘的图谱存盘Gexf文件，输出为用于Neo4j的导入数据文件
        :param filename:存盘graphml文件
        :return:Neo4j的导入数据文件
        '''
        input_filename = "{}//results//graph//{}".format(PROJECT_ROOT, filename)
        G = nx.read_graphml(input_filename)
        print("number of nodes: ", G.number_of_nodes(), ", number of dges: ", G.number_of_edges())

        '''
        >>> nodes.csv
        id:ID,Label,Weight,Type
        30073,D硬化性淋巴细胞性乳腺炎,1,D
        10032,B上皮,1,B

        >>> relations.csv
        :START_ID,:END_ID,Weight,:TYPE
        30073,10032,1,A
        30073,20037,1,A
        '''
        node_file = PROJECT_ROOT + "//results//graph//nodes.csv"
        with open(node_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['id:ID', 'name', 'Weight:float', 'label:LABEL'])
            for item in G.nodes(data=True):
                # item -> ('10017', {'type': 'B', 'name': '左乳', 'weight': 270})
                row = [int(item[0]), item[1]['name'], item[1]['weight'],item[1]['type']]
                # print(row)
                writer.writerow(row)

        edge_file = PROJECT_ROOT + "//results//graph//relations.csv"
        with open(edge_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow([':START_ID', ':END_ID', 'Weight:float', 'type:TYPE'])
            for item in G.edges(data=True):
                # item -> ('10017', '30006', {'weight': 58})
                row = [int(item[0]), int(item[1]), item[2]['weight'], "R"]
                writer.writerow(row)